from pathlib import Path
import httpx
import torch
import whisper
from bilibili_api import sync, video, Credential, HEADERS
import gpustat
from functional import seq

def get_mechineInfo():
    return [{
        'idx': i, 'temperature': card['temperature.gpu'], 
        'percent':card['utilization.gpu'], 'occupy':card['memory.used'], 
        'users': [
            {'name': p['username'], 'occupy': p['gpu_memory_usage']}
            for p in card['processes'] if p['gpu_memory_usage'] > 100
        ]}
        for i, card in enumerate(gpustat.new_query().jsonify()['gpus'])
    ]

def getIdle():
    idles = seq(get_mechineInfo()).filter(lambda x: x['occupy'] < 1000).map(lambda x: int(x['idx'])).list()
    assert len(idles) > 0
    return idles[0]

SESSDATA = ""
BILI_JCT = ""
BUVID3 = ""

# FFMPEG 路径，查看：http://ffmpeg.org/
FFMPEG_PATH = "ffmpeg"

async def download_url(url: str, out: str, info: str):
    # 下载函数
    async with httpx.AsyncClient(headers=HEADERS) as sess:
        resp = await sess.get(url)
        length = resp.headers.get('content-length')
        with open(out, 'wb') as f:
            process = 0
            for chunk in resp.iter_bytes(1024):
                if not chunk:
                    break

                process += len(chunk)
                f.write(chunk)

async def download(v: video.Video):
    # 获取视频下载链接
    download_url_data = await v.get_download_url(0)
    # 解析视频下载信息
    detecter = video.VideoDownloadURLDataDetecter(data=download_url_data)
    streams = detecter.detect_best_streams()
    # 有 MP4 流 / FLV 流两种可能
    temp_file = None
    if detecter.check_flv_stream() == True:
        # FLV 流下载
        temp_file = "./temp/temp.flv"
        await download_url(streams[0].url, temp_file, "FLV 音视频流")
        # # 转换文件格式
        # os.system(f'{FFMPEG_PATH} -i flv_temp.flv video.mp4')
        # # 删除临时文件
        # os.remove("flv_temp.flv")
        return temp_file
    else:
        # MP4 流下载
        # await download_url(streams[0].url, f"{target_name}.m4s", "视频流")
        temp_file = "./temp/temp.m4s"
        await download_url(streams[1].url, temp_file, "音频流")
        # # 混流
        # os.system(f'{FFMPEG_PATH} -i video_temp.m4s -i audio_temp.m4s -vcodec copy -acodec copy video.mp4')
        # # 删除临时文件
        # os.remove("video_temp.m4s")
        # os.remove("audio_temp.m4s")
        return temp_file

def stt(audio_path: str):
    model = whisper.load_model("medium", device=torch.device(f"cuda:{getIdle()}"))
    result = whisper.transcribe(model, audio_path, temperature=(0.0, 0.2, 0.4, 0.6000000000000001, 0.8, 1.0), **{'verbose': False, 'task': 'transcribe', 'language': 'zh', 'best_of': 5, 'beam_size': 5, 'patience': None, 'length_penalty': None, 'suppress_tokens': '-1', 'initial_prompt': None, 'condition_on_previous_text': True, 'fp16': True, 'compression_ratio_threshold': 2.4, 'logprob_threshold': -1.0, 'no_speech_threshold': 0.6, 'word_timestamps': False, 'prepend_punctuations': '"\'“¿([{-', 'append_punctuations': '"\'.。,，!！?？:：”)]}、'})
    return '\n'.join([
        seg['text']
        for seg in result['segments']
    ])
    # return result['text'] # 返回整段文本，没有换行

async def download_text(vid="BV16y411B7Rt"):
    credential = Credential(sessdata=SESSDATA, bili_jct=BILI_JCT, buvid3=BUVID3)
    v = video.Video(bvid=vid, credential=credential)
    video_path = await download(v)
    print("Downloaded video")
    info = await v.get_info()
    title = info["title"]
    text = stt(video_path)
    Path(f"{title}.txt").write_text(text)

if __name__ == '__main__':
    import sys
    sync(download_text(sys.argv[1]))
    
    
